import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


class PointWiseFeedForward(nn.Module):
    def __init__(self, hidden_units, dropout_rate):  # wried, why fusion X 2?

        super(PointWiseFeedForward, self).__init__()

        self.conv1 = nn.Conv1d(hidden_units, hidden_units, kernel_size=1)
        self.dropout1 = nn.Dropout(p=dropout_rate)
        self.conv2 = nn.Conv1d(hidden_units, hidden_units, kernel_size=1)
        self.dropout2 = nn.Dropout(p=dropout_rate)

    def forward(self, inputs):
        outputs = self.dropout2(self.conv2(
            self.dropout1(self.conv1(inputs.transpose(-1, -2)))))
        # as Conv1D requires (N, C, Length)
        outputs = outputs.transpose(-1, -2)
        outputs += inputs
        return outputs


class TimeAwareMultiHeadAttention(nn.Module):
    # required homebrewed mha layer for Ti/SASRec experiments
    def __init__(self, hidden_size, head_num, dropout_rate, dev):
        super(TimeAwareMultiHeadAttention, self).__init__()
        self.Q_w = nn.Linear(hidden_size, hidden_size)
        self.K_w = nn.Linear(hidden_size, hidden_size)
        self.V_w = nn.Linear(hidden_size, hidden_size)

        self.dropout = nn.Dropout(p=dropout_rate)
        self.softmax = nn.Softmax(dim=-1)

        self.hidden_size = hidden_size
        self.head_num = head_num
        self.head_size = hidden_size // head_num
        self.dropout_rate = dropout_rate
        self.dev = dev

    def forward(self, queries, keys, time_mask, attn_mask, abs_pos_K, abs_pos_V):
        Q, K, V = self.Q_w(queries), self.K_w(keys), self.V_w(keys)

        # head dim * batch dim for parallelization (h*N, T, C/h)
        Q_ = torch.cat(torch.split(Q, self.head_size, dim=2), dim=0)
        K_ = torch.cat(torch.split(K, self.head_size, dim=2), dim=0)
        V_ = torch.cat(torch.split(V, self.head_size, dim=2), dim=0)

        abs_pos_K_ = torch.cat(torch.split(
            abs_pos_K, self.head_size, dim=2), dim=0)
        abs_pos_V_ = torch.cat(torch.split(
            abs_pos_V, self.head_size, dim=2), dim=0)

        # batched channel wise matmul to gen attention weights
        attn_weights = Q_.matmul(torch.transpose(K_, 1, 2))
        attn_weights += Q_.matmul(torch.transpose(abs_pos_K_, 1, 2))

        # seq length adaptive scaling
        attn_weights = attn_weights / (K_.shape[-1] ** 0.5)

        # key masking, -2^32 lead to leaking, inf lead to nan
        # 0 * inf = nan, then reduce_sum([nan,...]) = nan

        time_mask = time_mask.unsqueeze(-1).expand(
            attn_weights.shape[0], -1, attn_weights.shape[-1])
        attn_mask = attn_mask.unsqueeze(0).expand(
            attn_weights.shape[0], -1, -1)
        paddings = torch.ones(attn_weights.shape) * \
            (-2**32+1)  # -1e23 # float('-inf')
        paddings = paddings.to(self.dev)
        attn_weights = torch.where(
            time_mask, paddings, attn_weights)  # True:pick padding
        attn_weights = torch.where(
            attn_mask, paddings, attn_weights)  # enforcing causality

        # code as below invalids pytorch backward rules
        attn_weights = self.softmax(attn_weights)
        # attn_weights = torch.where(time_mask, paddings, attn_weights) # weird query mask in tf impl
        # https://discuss.pytorch.org/t/how-to-set-nan-in-tensor-to-0/3918/4
        # attn_weights[attn_weights != attn_weights] = 0 # rm nan for -inf into softmax case
        attn_weights = self.dropout(attn_weights)
        outputs = attn_weights.matmul(V_)
        outputs += attn_weights.matmul(abs_pos_V_)

        # (num_head * N, T, C / num_head) -> (N, T, C)
        outputs = torch.cat(torch.split(
            outputs, Q.shape[0], dim=0), dim=2)  # div batch_size

        return outputs


class SASRec(nn.Module):  # similar to nn.MultiheadAttention
    def __init__(self, hyper_params):
        super(SASRec, self).__init__()

        self.hyper_params = hyper_params

        self.user_num = hyper_params['total_users']
        self.item_num = hyper_params['total_items']
        self.dev = hyper_params['device']

        # TODO: loss += args.l2_emb for regularizing embedding vectors during training
        # https://stackoverflow.com/questions/42704283/adding-l1-l2-regularization-in-pytorch
        self.item_emb = nn.Embedding(
            self.item_num+1, hyper_params['hidden_size'], padding_idx=0)
        self.item_emb_dropout = nn.Dropout(p=hyper_params['dropout_rate'])

        self.abs_pos_K_emb = nn.Embedding(
            hyper_params['seq_len'], hyper_params['hidden_size'])
        self.abs_pos_V_emb = nn.Embedding(
            hyper_params['seq_len'], hyper_params['hidden_size'])
        self.time_matrix_K_emb = nn.Embedding(
            self.item_num+1, hyper_params['hidden_size'])
        self.time_matrix_V_emb = nn.Embedding(
            hyper_params['time_span'] + 1, hyper_params['hidden_size'])

        self.item_emb_dropout = nn.Dropout(p=hyper_params['dropout_rate'])
        self.abs_pos_K_emb_dropout = nn.Dropout(p=hyper_params['dropout_rate'])
        self.abs_pos_V_emb_dropout = nn.Dropout(p=hyper_params['dropout_rate'])
        self.time_matrix_K_dropout = nn.Dropout(p=hyper_params['dropout_rate'])
        self.time_matrix_V_dropout = nn.Dropout(p=hyper_params['dropout_rate'])

        self.attention_layernorms = nn.ModuleList()  # to be Q for self-attention
        self.attention_layers = nn.ModuleList()
        self.forward_layernorms = nn.ModuleList()
        self.forward_layers = nn.ModuleList()

        self.last_layernorm = nn.LayerNorm(
            hyper_params['hidden_size'], eps=1e-8)

        for _ in range(hyper_params['num_blocks']):
            new_attn_layernorm = nn.LayerNorm(
                hyper_params['hidden_size'], eps=1e-8)
            self.attention_layernorms.append(new_attn_layernorm)

            new_attn_layer = TimeAwareMultiHeadAttention(hyper_params['hidden_size'],
                                                         hyper_params['num_heads'],
                                                         hyper_params['dropout_rate'],
                                                         hyper_params['device'])
            self.attention_layers.append(new_attn_layer)

            new_fwd_layernorm = nn.LayerNorm(
                hyper_params['hidden_size'], eps=1e-8)
            self.forward_layernorms.append(new_fwd_layernorm)

            new_fwd_layer = PointWiseFeedForward(
                hyper_params['hidden_size'], hyper_params['dropout_rate'])
            self.forward_layers.append(new_fwd_layer)

        self.b_o = nn.Parameter(torch.Tensor(
            hyper_params['total_items'] + 1), requires_grad=True)
        self.register_parameter('b_o', self.b_o)
        torch.nn.init.normal_(self.b_o)
        # self.pos_sigmoid = nn.Sigmoid()
        # self.neg_sigmoid = nn.Sigmoid()

    def seq2feats(self, log_seqs):
        seqs = self.item_emb(log_seqs)
        seqs *= self.item_emb.embedding_dim ** 0.5
        seqs = self.item_emb_dropout(seqs)

        positions = np.tile(np.array(range(log_seqs.shape[1])), [
                            log_seqs.shape[0], 1])
        positions = torch.LongTensor(positions).to(self.dev)
        abs_pos_K = self.abs_pos_K_emb(positions)
        abs_pos_V = self.abs_pos_V_emb(positions)
        abs_pos_K = self.abs_pos_K_emb_dropout(abs_pos_K)
        abs_pos_V = self.abs_pos_V_emb_dropout(abs_pos_V)

        # mask 0th items(placeholder for dry-run) in log_seqs
        # would be easier if 0th item could be an exception for training
        timeline_mask = log_seqs == 0
        seqs *= ~timeline_mask.unsqueeze(-1)  # broadcast in last dim

        tl = seqs.shape[1]  # time dim len for enforce causality
        attention_mask = ~torch.tril(torch.ones(
            (tl, tl), dtype=torch.bool, device=self.dev))

        for i in range(len(self.attention_layers)):
            # Self-attention, Q=layernorm(seqs), K=V=seqs
            # seqs = torch.transpose(seqs, 0, 1) # (N, T, C) -> (T, N, C)
            # PyTorch mha requires time first fmt
            Q = self.attention_layernorms[i](seqs)
            mha_outputs = self.attention_layers[i](Q, seqs,
                                                   timeline_mask, attention_mask,
                                                   abs_pos_K, abs_pos_V)
            seqs = Q + mha_outputs
            # seqs = torch.transpose(seqs, 0, 1) # (T, N, C) -> (N, T, C)

            # Point-wise Feed-forward, actually 2 Conv1D for channel wise fusion
            seqs = self.forward_layernorms[i](seqs)
            seqs = self.forward_layers[i](seqs)
            seqs *= ~timeline_mask.unsqueeze(-1)

        log_feats = self.last_layernorm(seqs)

        return log_feats

    def forward(self, log_seqs):  # for training
        log_feats = self.seq2feats(log_seqs)

        pred = torch.matmul(log_feats, self.item_emb.weight.t()) + self.b_o
        pred = torch.softmax(pred, dim=-1)

        return pred  # pos_pred, neg_pred

    def predict(self, log_seqs):  # for inference
        log_feats = self.seq2feats(log_seqs)

        # only use last QKV classifier, a waste
        final_feat = log_feats[:, -1, :]

        pred = torch.matmul(final_feat, self.item_emb.weight.t()) + self.b_o
        pred = torch.softmax(pred, dim=-1)

        return pred  # preds # (U, I)


class Model(nn.Module):
    def __init__(self, hyper_params):
        super(Model, self).__init__()

        self.hyper_params = hyper_params
        self.embedding = nn.Embedding(
            num_embeddings=hyper_params['total_items'] + 1,
            embedding_dim=hyper_params['item_embed_size'],
            padding_idx=0
        )
        self.linear_q = nn.Linear(
            hyper_params['item_embed_size'], hyper_params['hidden_size'])
        self.linear_k = nn.Linear(
            hyper_params['item_embed_size'], hyper_params['hidden_size'])
        self.linear_v = nn.Linear(
            hyper_params['item_embed_size'], hyper_params['hidden_size'])

        self.p_e = nn.Embedding(
            num_embeddings=hyper_params['seq_len'], embedding_dim=hyper_params['item_embed_size'])

        self.attn = nn.MultiheadAttention(
            embed_dim=hyper_params['hidden_size'],
            num_heads=1,
            dropout=hyper_params['dropout_rate']
        )
        self.transformer_layer = nn.TransformerEncoderLayer(
            d_model=hyper_params['item_embed_size'],
            nhead=2,
            dim_feedforward=hyper_params['hidden_size'],
            dropout=hyper_params['dropout_rate'],
            activation='gelu'
        )
        self.transformer = nn.TransformerEncoder(
            encoder_layer=self.transformer_layer,
            num_layers=2
        )

        self.linear_o = nn.Linear(
            hyper_params['hidden_size'], hyper_params['total_items']+1)
        self.dropout = nn.Dropout()

    def forward(self, x, padding):
        positions = np.tile(np.array(range(x.shape[1])), [x.shape[0], 1])
        positions = torch.LongTensor(positions).to(self.hyper_params['device'])
        embed = self.embedding(x)
        embed = embed + self.p_e(positions)

        embed = self.dropout(embed)
        transformer_out = self.transformer(
            src=embed.transpose(0, 1),
            src_key_padding_mask=padding
        )

        pred = self.linear_o(transformer_out.transpose(0, 1))
        pred = torch.softmax(pred, dim=-1)

        return pred
